// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::{hash_map::Entry, HashMap}, io::Write, path::Path};

use anyhow::Context;
use log::info;

use crate::{database::deploy_stat::{DeployStat, DeployStats}, deploy::deploy::deploy_pkgs_for_rollback, Arch};

use super::{undeploy_pkgs, MergeMode};

type RollbackPkgs = HashMap<Arch, HashMap<String, (Option<MergeMode>, Option<MergeMode>)>>;

fn rollback_forward(current_stats: &[DeployStat], extra_stats: &[DeployStat]) -> RollbackPkgs {
    let mut to_operate: RollbackPkgs = HashMap::new();

    for stat in current_stats {
        match to_operate.entry(stat.arch) {
            Entry::Occupied(mut occupied_entry) => {
                for pkg in &stat.pkgs {
                    occupied_entry.get_mut().insert(pkg.to_string(), (stat.mode, stat.mode));
                }
            },
            Entry::Vacant(vacant_entry) => {
                let mut pkgs = HashMap::new();
                for pkg in &stat.pkgs {
                    pkgs.insert(pkg.to_string(), (stat.mode, stat.mode));
                }
                vacant_entry.insert(pkgs);
            },
        }
    }

    for stat in extra_stats {
        for pkg in &stat.pkgs {
            match to_operate.entry(stat.arch) {
                Entry::Occupied(mut occupied_entry) => {
                    match occupied_entry.get_mut().entry(pkg.to_string()) {
                        Entry::Occupied(mut occupied_entry) => {
                            occupied_entry.get_mut().1 = stat.mode;
                        },
                        Entry::Vacant(vacant_entry) => {
                            vacant_entry.insert((None, stat.mode));
                        },
                    }
                },
                Entry::Vacant(vacant_entry) => {
                    let mut pkgs = HashMap::new();
                    pkgs.insert(pkg.to_string(), (None, stat.mode));
                    vacant_entry.insert(pkgs);
                },
            }
        }
    }
    to_operate
}

fn rollback_backward(target_stats: &[DeployStat], extra_stats: &[DeployStat]) -> RollbackPkgs {
    let mut to_operate: RollbackPkgs = HashMap::new();

    for stat in target_stats {
        for pkg in &stat.pkgs {
            match to_operate.entry(stat.arch) {
                Entry::Occupied(mut occupied_entry) => {
                    occupied_entry.get_mut().insert(pkg.to_string(), (stat.mode, stat.mode));
                },
                Entry::Vacant(vacant_entry) => {
                    let mut pkgs = HashMap::new();
                    pkgs.insert(pkg.to_string(), (stat.mode, stat.mode));
                    vacant_entry.insert(pkgs);
                },
            }
        }
    }

    for stat in extra_stats {
        for pkg in &stat.pkgs {
            match to_operate.entry(stat.arch) {
                Entry::Occupied(mut occupied_entry) => {
                    match occupied_entry.get_mut().entry(pkg.to_string()) {
                        Entry::Occupied(mut occupied_entry) => {
                            occupied_entry.get_mut().0 = stat.mode;
                        },
                        Entry::Vacant(vacant_entry) => {
                            vacant_entry.insert((stat.mode, None));
                        },
                    }
                },
                Entry::Vacant(vacant_entry) => {
                    let mut pkgs = HashMap::new();
                    pkgs.insert(pkg.to_string(), (stat.mode, None));
                    vacant_entry.insert(pkgs);
                },
            }
        }
    }
    to_operate
}

pub async fn collect_rollback_pkgs<P: AsRef<Path>>(
    dir: P,
    current_stats: &DeployStats,
    current_id: usize,
    target_id: usize,
) -> Result<RollbackPkgs, anyhow::Error> {
    if target_id > current_id {
        let extra_stats = DeployStat::load_id_range(dir.as_ref(), (current_id+1)..=target_id).await
            .context(format!("failed to load extra deploy stats in '{}'", dir.as_ref().display()))?;
        Ok(rollback_forward(&current_stats.stats, &extra_stats))
    } else if target_id < current_id {
        let extra_stats = &current_stats.stats[target_id..];
        let target_stats = &current_stats.stats[..target_id];
        Ok(rollback_backward(target_stats, extra_stats))
    } else { 
        Err(anyhow::anyhow!("input id is same with current id"))
    }
}

fn need_undeploy(init_mode: &Option<MergeMode>, new_mode: &Option<MergeMode>) -> bool {
    match (init_mode, new_mode) {
        (Some(_), None) => true, 
        (Some(init_mode), Some(new_mode)) => init_mode != new_mode,
        _ => false,
    }
}

fn match_deploy_mode(init_mode: &Option<MergeMode>, new_mode: &Option<MergeMode>) -> Option<MergeMode> {
    match (init_mode, new_mode) {
        (None, Some(new_mode)) => Some(*new_mode),
        (Some(init_mode), Some(new_mode)) => {
            if init_mode != new_mode {
                Some(*new_mode)
            } else {
                None
            }
        },
        _ => None,
    }
}

pub async fn rollback<P: AsRef<Path>>(
    root: P,
    current_stats: &DeployStats,
    rollback_pkgs: RollbackPkgs,
) -> anyhow::Result<()> {
    {
        let mut dump = Vec::new();
        if dump_rollback_pkgs(&rollback_pkgs, &mut dump).is_ok() {
            info!("\n{}", String::from_utf8_lossy(&dump));
        }
    }

    for (arch, pkgs) in &rollback_pkgs {
        let to_undeploy_pkgs = pkgs
            .iter()
            .filter(|(_, (init_mode, new_mode))| need_undeploy(init_mode, new_mode))
            .map(|(pkg, _)| pkg.as_str());
        undeploy_pkgs(current_stats, root.as_ref(), to_undeploy_pkgs, *arch).await
            .context("failed to undeploy pkgs")?;
    }

    for (arch, pkgs) in &rollback_pkgs {
        let to_deploy_pkgs = pkgs
            .iter()
            .filter_map(|(pkg, (init_mode, new_mode))|
                match_deploy_mode(init_mode, new_mode).map(|mode| (pkg.as_str(), mode)));
        deploy_pkgs_for_rollback(root.as_ref(), to_deploy_pkgs, arch).await?;
    }

    Ok(())
}

fn dump_rollback_pkgs<W: Write>(
    rollback_pkgs: &RollbackPkgs,
    mut w: W,
) -> std::io::Result<()> {
    writeln!(w, "Undeploy pkgs:")?;
    for (arch, pkgs) in rollback_pkgs {
        for (pkg, (init_mode, new_mode)) in pkgs {
            if need_undeploy(init_mode, new_mode) {
                writeln!(w, "  - {}, deploy arch: {}", pkg, arch)?;
            }
        }
    }
    writeln!(w, "Deploy pkgs:")?;
    for (arch, pkgs) in rollback_pkgs {
        for (pkg, (init_mode, new_mode)) in pkgs {
            if let Some(mode) = match_deploy_mode(init_mode, new_mode) {
                writeln!(w, "  - {}: {}, deploy arch: {}", pkg, mode.to_string(), arch)?;
            }
        }
    }
    Ok(())
}
